# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: chunhui

import json
import os
import shutil
import subprocess
import warnings
import tempfile
import time
import glob

import numpy as np
import pandas as pd

from cc.appcoin2.strategy.support.util import get_cmd_vault_key_copy
from python.coin_deploy.basis_strat2.script.adjust_universe import get_max_pos
from coin.exchange.okex_futures_v5.kr_rest.native_private_client import \
    OkexFuturesNativePrivateClient

from absl import app, flags


FLAGS = flags.FLAGS

def _get_okex_product_info():
  with open('data/coin2/product_info/Futures.Okex.v5-swap.json') as f:
    product_info = json.load(f)
  return product_info


class OkexSwapClient(OkexFuturesNativePrivateClient):
  def __init__(self, key_file):
    super().__init__(key_file=key_file)
    pi = _get_okex_product_info()
    self.mea = pi['mea']
    self.product_info = pi['product_infos']

  def convert_native_symbol(self, symbol):
    for info in self.product_info:
      if info['symbol'] == symbol:
        return info['native_symbol']
    raise ValueError(f'Unrecognized symbol {symbol}')

  def _cancel_orders(self, symbol, cancel_ord_id_list):
    method = 'POST'
    path = '/api/v5/trade/cancel-order'
    native_symbol = self.convert_native_symbol(symbol)

    for ord_id in cancel_ord_id_list:
      params = {
          'instId': native_symbol,
          'ordId': ord_id
      }
      update = self._query(method, path, params)
      print(update)

  def cancel_orders(self, symbol, pd_orders):
    result = self._cancel_orders(symbol, pd_orders)
    print(result)

  def _get_pending_orders(self, symbol):
    method = 'GET'
    path = '/api/v5/trade/orders-pending'
    native_symbol = self.convert_native_symbol(symbol)
    params = {
        'instId': native_symbol,
    }
    update = self._query(method, path, params)
    return update

  def get_pending_orders(self, symbol):
    result = self._get_pending_orders(symbol)
    return [ele['ordId'] for ele in result.as_json()['msg']['data']]

  def _set_leverage(self, symbol, leverage):
    method = 'POST'
    path = '/api/v5/account/set-leverage'
    native_symbol = self.convert_native_symbol(symbol)
    params = {
        'instId': native_symbol,
        'lever': leverage,
        'mgnMode': 'cross',
        'posSide': 'net'
    }
    update = self._query(method, path, params)
    return update

  def set_leverage(self, symbols, leverage):
    if isinstance(symbols, str):
      symbols = symbols.split(',')
    for i, symbol in enumerate(symbols):
      time.sleep(0.11)
      result = self._set_leverage(symbol, leverage)
      print(result)

  def _get_pos_limit(self, symbol):
    # https://www.okx.com/docs-v5/en/#rest-api-public-data-get-position-tiers
    method = 'GET'
    path = '/api/v5/public/position-tiers'
    native_symbol = self.convert_native_symbol(symbol)
    instFamily = native_symbol.split('-')[0] + '-USDT'
    params = {
        'instType': 'SWAP',
        "instFamily": instFamily,
        "tdMode": "cross",
        "ccy": "USDT"
    }
    update = self._query(method, path, params)
    return update

  def _get_leverage(self, symbol):
    method = 'GET'
    path = '/api/v5/account/leverage-info/'
    native_symbol = self.convert_native_symbol(symbol)
    params = {
        "instId": native_symbol,
        "mgnMode": "cross"
    }
    update = self._query(method, path, params)
    return update

  def get_leverage(self, symbols):
    if isinstance(symbols, str):
      symbols = symbols.split(',')
    data = []
    for i, symbol in enumerate(symbols):
      if (i + 1) % 19 == 0:
        time.sleep(2)
      leverage = self._get_leverage(symbol).as_json()['msg']['data']
      pos_limit = self._get_pos_limit(symbol)
      data.extend(leverage)
      df = pd.DataFrame(pos_limit.as_json()['msg']['data'])
      df['maxLever'] = df['maxLever'].astype(float)
      df['norm_symbol'] = symbol
      sorted_df = df.sort_values('maxLever', ascending=True)
      value = float(leverage[0]['lever'])
      row = sorted_df[sorted_df['maxLever'] <= value].iloc[-1]
      data[-1].update(row.to_dict())
    df = pd.DataFrame(data)
    df['symbol'] = df['instFamily'].apply(lambda x: x.split('-')[0])
    price_files = glob.glob(f'{FLAGS.price_cache_root}/okex_swap_prices*.json')
    if len(price_files) > 0:
      price_json = json.load(open(max(price_files, key=os.path.getctime)))
      price_dict = {info['instId'].split('-')[0]: float(info['last'])
                                            for info in price_json['data']}
      multipliers = {info['native_symbol'].split('-')[0]: info['contract_value']
                                          for info in self.product_info}
      df['price'] = df['symbol'].map(price_dict)
      df['multipliers'] = df['symbol'].map(multipliers)
      df['maxSzDollar'] = df['maxSz'].astype(float) * df['price'] * df['multipliers']
    else:
      df['price'] = df['multipliers'] = df['maxSzDollar'] = np.nan
    col = "symbol,norm_symbol,lever,mgnMode,posSide,maxLever,price,maxSz,maxSzDollar".split(',')
    return df[col].set_index('symbol')


def init_flags():
  flags.DEFINE_string('strategy', '', 'strategy_name')
  flags.DEFINE_string('account', '', 'account name')
  flags.DEFINE_string('symbol', '', 'usdt amount')
  flags.DEFINE_integer('leverage', 5, 'upd leverage')
  flags.DEFINE_string('key_filepath', '', 'key filepath')
  flags.DEFINE_bool('query_only', True, 'check leverage')
  flags.DEFINE_bool('update_all', False, 'update all symbols')
  flags.DEFINE_string('price_cache_root',
                      '/remote/iosg/home-3/yuxuan/shared/prices/', '')


def get_okex_account(strategy):
  if os.path.exists(f'python/coin_deploy/basis_strat2/driver/{strategy}.json'):
    driver = json.load(open(f'python/coin_deploy/basis_strat2/driver/{strategy}.json'))
    for _, config in driver['order'].items():
      if config['mea'] == 'Futures.Okex.v5-swap':
        return config['connection_config']['key_filepath'].split('/')[3], config['products']
  raise ValueError(f"look up account for {strategy} failed")


def check_max_pos(msc, df, strat_symbols):
  if not flags.FLAGS.strategy:
    return
  max_pos = get_max_pos(flags.FLAGS.strategy)
  for symbol in strat_symbols:
    max_notional_value = df.loc[msc.convert_native_symbol(symbol).split('-')[0]]['maxSzDollar']
    if max_notional_value and max_notional_value < max_pos:
      warnings.warn(
          f'{flags.FLAGS.strategy} {symbol} max_allowed_postion {max_notional_value} is smaller than the current max position {max_pos}')


def main(_):
  pd.set_option('display.max_rows', None)
  strat_symbols = list()
  if flags.FLAGS.strategy:
    account, strat_symbols = get_okex_account(flags.FLAGS.strategy)
  elif flags.FLAGS.key_filepath:
    account = "use custom keypath"
  else:
    assert flags.FLAGS.account, "None Account Id"
    account = flags.FLAGS.account
  if flags.FLAGS.update_all:
    all_products = _get_okex_product_info()
    strat_symbols = [s['symbol'] for s in all_products['product_infos']]
  elif flags.FLAGS.symbol:
    strat_symbols = flags.FLAGS.symbol.split(',')
  assert len(strat_symbols) > 0, "No Strat Symbol"
  print(f'{flags.FLAGS.strategy} is using account {account}')
  tmp_dir = tempfile.mkdtemp()

  try:
    if flags.FLAGS.key_filepath:
      key_file = flags.FLAGS.key_filepath
    else:
      key_file = os.path.join(tmp_dir, 'key.json')
      vault_key_name = f"trade/{account}/OKEx/trade_v5_key.json"
      cmd = get_cmd_vault_key_copy(vault_key_name, key_file)
      _ = subprocess.check_output(cmd, shell=True)
    msc = OkexSwapClient(key_file=key_file)
    result = msc.get_leverage(strat_symbols)
    print("\n#####################")
    print("-- Before Adjust --")
    print(result)
    print("######################\n")

    if not flags.FLAGS.query_only:
      for symbol in strat_symbols:
        curr_lever = result[result['norm_symbol'] == symbol]['lever'].astype(int).tolist()
        if len(curr_lever) != 1:
          print(f"{symbol} wrong lever info.. skip")
          continue
        if FLAGS.leverage == curr_lever[0]:
          print(f"{symbol} Already applied. skip.. leverage = {curr_lever[0]}")
          continue
        pd_orders = msc.get_pending_orders(symbol)
        msc.cancel_orders(symbol, pd_orders)
        msc.set_leverage([symbol], FLAGS.leverage)
      print("\n#####################")
      print("-- After Adjust --")
      df = msc.get_leverage(strat_symbols)
      print(df)
      print("######################\n")
      check_max_pos(msc, df, strat_symbols)
  except Exception:
    raise
  finally:
    shutil.rmtree(tmp_dir)


'''
./pyrunner cc/appcoin2/strategy/support/okex_leverage.py \
--strategy=model_smm_okex_perpetual_6 \
--query_only=1 \
--leverage=5 \
--symbol=OP-USDT.PERPETUAL
'''

'''
./pyrunner cc/appcoin2/strategy/support/okex_leverage.py \
--account=prestoinvt24 \
--query_only=0 --leverage=5 \
--symbol=CONV-USDT.PERPETUAL,TORN-USDT.PERPETUAL
'''

if __name__ == '__main__':
  init_flags()
  app.run(main)
